#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RAG智能对话机器人 - Web界面版本
基于Flask提供Web聊天界面
"""

import json
from flask import Flask, render_template, request, jsonify, Response
from datetime import datetime
from tools.RAGChatBot import RAGChatBot

app = Flask(__name__)
app.secret_key = 'rag_chatbot_secret_key'

# 全局RAG机器人实例
bot = None


def initialize_rag_bot():
    """初始化RAG机器人"""
    global bot
    try:
        bot = RAGChatBot(
            ollama_url="http://172.26.32.1:11434",
            milvus_collection="workspace_files"
        )
        
        if bot.initialize():
            print("✅ RAG系统初始化成功")
            return True
        else:
            print("❌ RAG系统初始化失败")
            return False
    except Exception as e:
        print(f"❌ 初始化RAG机器人失败: {e}")
        return False


@app.route('/')
def index():
    """主页"""
    return render_template('chat.html')


@app.route('/api/chat', methods=['POST'])
def chat():
    """处理聊天请求"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        data = request.get_json()
        user_message = data.get('message', '').strip()
        
        if not user_message:
            return jsonify({
                'success': False,
                'error': '消息不能为空'
            })
        
        # 生成回答
        response, sources = bot.chat(user_message)
        
        # 格式化源文件信息
        formatted_sources = []
        if sources:
            for source in sources:
                formatted_sources.append({
                    'file_name': source['file_name'],
                    'file_path': source['file_path'],
                    'file_type': source['file_type'],
                    'similarity_score': round(source['similarity_score'], 3),
                    'modified_time': source['modified_time']
                })
        
        return jsonify({
            'success': True,
            'response': response,
            'sources': formatted_sources,
            'timestamp': datetime.now().isoformat()
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/api/chat/stream', methods=['POST'])
def chat_stream():
    """流式聊天响应"""
    if not bot:
        return Response(
            json.dumps({'error': 'RAG系统未初始化'}),
            mimetype='application/json'
        )
    
    try:
        data = request.get_json()
        user_message = data.get('message', '').strip()
        
        if not user_message:
            return Response(
                json.dumps({'error': '消息不能为空'}),
                mimetype='application/json'
            )
        
        def generate():
            try:
                # 首先搜索文件
                retrieved_files = bot.search_files(user_message)
                
                # 发送源文件信息
                if retrieved_files:
                    sources_data = {
                        'type': 'sources',
                        'sources': [{
                            'file_name': f['file_name'],
                            'file_path': f['file_path'],
                            'similarity_score': round(f['similarity_score'], 3)
                        } for f in retrieved_files]
                    }
                    yield f"data: {json.dumps(sources_data)}\n\n"
                
                # 构建上下文
                context = bot.build_context(user_message, retrieved_files)
                
                # 流式生成回答
                response_text = ""
                for chunk in bot.generate_response_stream(user_message, context):
                    response_text += chunk
                    chunk_data = {
                        'type': 'chunk',
                        'content': chunk
                    }
                    yield f"data: {json.dumps(chunk_data)}\n\n"
                
                # 发送完成信号
                final_data = {
                    'type': 'done',
                    'full_response': response_text
                }
                yield f"data: {json.dumps(final_data)}\n\n"
                
            except Exception as e:
                error_data = {
                    'type': 'error',
                    'error': str(e)
                }
                yield f"data: {json.dumps(error_data)}\n\n"
        
        return Response(
            generate(),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'Connection': 'keep-alive',
                'Access-Control-Allow-Origin': '*'
            }
        )
        
    except Exception as e:
        return Response(
            json.dumps({'error': str(e)}),
            mimetype='application/json'
        )


@app.route('/api/models', methods=['GET'])
def get_models():
    """获取可用模型"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        models = bot.get_available_models()
        return jsonify({
            'success': True,
            'models': models,
            'current_model': bot.current_model
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/api/model', methods=['POST'])
def switch_model():
    """切换模型"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        data = request.get_json()
        model_name = data.get('model', '').strip()
        
        if not model_name:
            return jsonify({
                'success': False,
                'error': '模型名称不能为空'
            })
        
        success = bot.switch_model(model_name)
        return jsonify({
            'success': success,
            'current_model': bot.current_model
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/api/history', methods=['GET'])
def get_history():
    """获取对话历史"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        history = bot.get_chat_history()
        formatted_history = []
        
        for msg in history:
            formatted_msg = {
                'role': msg.role,
                'content': msg.content,
                'timestamp': datetime.fromtimestamp(msg.timestamp).isoformat(),
                'sources': msg.sources if hasattr(msg, 'sources') and msg.sources else []
            }
            formatted_history.append(formatted_msg)
        
        return jsonify({
            'success': True,
            'history': formatted_history
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/api/history', methods=['DELETE'])
def clear_history():
    """清空对话历史"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        bot.clear_history()
        return jsonify({
            'success': True,
            'message': '对话历史已清空'
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/api/stats', methods=['GET'])
def get_stats():
    """获取系统统计"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        stats = bot.get_stats()
        return jsonify({
            'success': True,
            'stats': stats
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


@app.route('/api/search', methods=['POST'])
def search_files():
    """搜索文件"""
    if not bot:
        return jsonify({
            'success': False,
            'error': 'RAG系统未初始化'
        })
    
    try:
        data = request.get_json()
        query = data.get('query', '').strip()
        top_k = data.get('top_k', 10)
        
        if not query:
            return jsonify({
                'success': False,
                'error': '搜索关键词不能为空'
            })
        
        results = bot.search_files(query, top_k=top_k)
        
        formatted_results = []
        for result in results:
            formatted_results.append({
                'file_name': result['file_name'],
                'file_path': result['file_path'],
                'file_type': result['file_type'],
                'file_size': result['file_size'],
                'modified_time': result['modified_time'],
                'similarity_score': round(result['similarity_score'], 3)
            })
        
        return jsonify({
            'success': True,
            'results': formatted_results,
            'total': len(formatted_results)
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        })


if __name__ == '__main__':
    print("🚀 启动RAG Web对话机器人...")
    
    if initialize_rag_bot():
        print("🌐 启动Web服务器...")
        print("访问地址: http://localhost:5000")
        app.run(host='0.0.0.0', port=5000, debug=True)
    else:
        print("❌ 初始化失败，程序退出")